#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/12/16 15:18
# @Author  : 新猿异码
# @Site    : https://gitee.com/chshao/autotest
# @CSDN    : https://blog.csdn.net/m0_37576542?type=blog
# @File    : adb.py
# @Desc    : 模块功能描述
# @Software: PyCharm
# -*- coding:utf-8 -*-
import os
import re
import time
import socket
import subprocess

import whichcraft
from collections import namedtuple
from autotest.buildin.PopenOS import Popen


class AdbError(Exception):
    def __init__(self, error):
        super(AdbError, self).__init__()
        self.errorMsg = error

    def __str__(self):
        return self.errorMsg


_OKAY = "OKAY"
_FAIL = "FAIL"
_DENT = "DENT"  # Directory Entity
_DONE = "DONE"
ForwardItem = namedtuple("ForwardItem", ["serial", "local", "remote"])


def where_adb():
    adb_path = whichcraft.which('adb')
    if adb_path is None:
        raise EnvironmentError("Can't find adb,please install adb first.")
    return adb_path


class _AdbStreamConnect(object):
    """
    连接adb服务
    即通过adb start-server在PC侧ANDROID_ADB_SERVER_PORT端口起的服务
    具体可以发送哪些命令需要查看adb协议
    """

    def __init__(self, host=None, port=None):
        self.__host = host
        self.__port = port
        self.__conn = self._connect()

    @property
    def conn(self):
        return self.__conn

    def _connect(self):
        adb_host = self.__host or os.environ.get('ANDROID_ADB_SERVER_HOST', '127.0.0.1')
        adb_port = self.__port or int(os.environ.get('ANDROID_ADB_SERVER_PORT', 7305))
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(5)
        for retry in range(2):
            if s.connect_ex((adb_host, adb_port)) == 0:
                return s
            # Log.error(f"cannot connect to daemon at tcp:{adb_port}: cannot connect to {adb_host}:{adb_port}")
            # Log.warn("please ensure adb server started?")
            subprocess.Popen('adb kill-server')
            subprocess.Popen(f'adb -P {adb_port} start-server')
            time.sleep(2)
        raise ConnectionRefusedError(f"cannot connect to {adb_host}:{adb_port}")

    def close(self):
        self.conn.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def send(self, cmd):
        # cmd: str表示期望cmd是字符串类型
        if not isinstance(cmd, str):
            cmd = str(cmd)
        self.conn.send("{:04x}{}".format(len(cmd), cmd).encode("utf-8"))

    def read(self, n):
        # -> str表示接口返回值为字符串
        return self.conn.recv(n).decode()

    def read_string(self):
        size = int(self.read(4), 16)
        return self.read(size)

    def read_until_close(self):
        content = ""
        while True:
            chunk = self.read(4096)
            if not chunk:
                break
            content += chunk
        return content

    def check_okay(self):
        # 前四位是状态码
        data = self.read(4)
        if data == _FAIL:
            raise AdbError(self.read_string())
        elif data == _OKAY:
            return
        raise AdbError("Unknown data: %s" % data)


class AdbClient(object):

    def connect_tcp(self, host, port):
        """
        PC跟Android设备通过TCP通信
        :param host: 设备IP
        :param port: 设备端口
        :return:
        """
        serial = ":".join((host, str(port)))
        with self._connect() as c:
            # TCP/IP方式连接，必须保证已经切换至TCP/IP连接方式
            Popen(f"adb -s {serial} tcpip {serial.split(':')[1]}").wait_until_end()
            p = Popen(f"adb -s {serial} connect {serial}", daemonic=True)
            if p.wait_until(f"connected to {serial}", timeout=5) is False:
                p.terminate()
                raise RuntimeError(f"can not connect {serial}, please check.")
            time.sleep(2)
            return True

    def server_version(self):
        """
        @summary:获取adb版本号
        @return :版本号1.0.41中的41
        """
        with self._connect() as c:
            c.send("host:version")
            c.check_okay()
            return int(c.read_string(), 16)

    def _connect(self):
        return _AdbStreamConnect()

    def forward(self, serial, local, remote, norebind=False):
        """
        @summary:给adb服务端发送host-serial:<sn>:forward:tcp:<pc_port>;tcp:<phone_port>进行端口转发
        @param serial:手机sn号,sn为None时按默认连接一部手机处理
        @param local:PC侧socket客户端端口
        @param remote:手机侧socket服务端端口
        @param norebind:fail if already forwarded when set to true
        @attention :PC跟手机通过USB方式通信
        """
        with self._connect() as c:
            cmds = ["host", "forward"]
            if serial:
                cmds = ["host-serial", serial, "forward"]
            if norebind:
                cmds.append("norebind")
            cmds.append("tcp:%s;tcp:%s" % (local, remote))
            c.send(":".join(cmds))
            c.check_okay()

    def forward_list(self, serial=None):
        """
        @summary:查看端口转发是否成功
        @param serial:手机sn号
        @attention :PC跟手机通过USB方式通信
        """
        with self._connect() as c:
            list_cmd = "host:list-forward"
            if serial:
                list_cmd = "host-serial:{}:list-forward".format(serial)
            c.send(list_cmd)
            c.check_okay()
            content = c.read_string()
            for line in content.splitlines():
                parts = line.split()
                if len(parts) != 3:
                    continue
                if serial and parts[0] != serial:
                    continue
                yield ForwardItem(*parts)

    def shell(self, serial, command):
        """
        @summary:执行shell命令
        @param serial:手机sn号
        @param command:要执行的命令
        @attention :只能执行adb shell命令
        """
        with self._connect() as c:
            c.send("host:transport:" + serial)
            c.check_okay()
            c.send("shell:" + command)
            c.check_okay()
            return c.read_until_close()

    def push(self, local, remote, serial=None) -> str:
        """
        将本地的文件资源导入手机
        :param local: 本地文件资源
        :param remote: 导入到手机的目标路径
        :param serial: 手机sn号
        :return:
        """
        if serial is None:
            self.must_one_device(serial)
            serial = self.device_list()[0]
        p = os.popen(f"adb -s {serial} push {local} {remote}")
        return p.read()

    def pull(self, remote, local, serial=None) -> str:
        """
        从手机导出文件资源
        :param remote: 手机上的文件资源
        :param local: 导出的目标路径
        :param serial: 手机sn号
        :return:
        """
        if serial is None:
            self.must_one_device(serial)
            serial = self.device_list()[0]
        p = os.popen(f"adb -s {serial} pull {remote} {local}")
        return p.read()

    def device_list(self):
        """
        @summary:获取手机列表
        @attention :
        """
        device_list = []
        with self._connect() as c:
            c.send("host:devices")
            c.check_okay()
            output = c.read_string()
            for line in output.splitlines():
                parts = line.strip().split("\t")
                if len(parts) != 2:
                    continue
                if parts[1] == 'device':
                    device_list.append(parts[0])
        return device_list

    def must_one_device(self, serial):
        device_list = self.device_list()
        if len(device_list) == 0:
            raise RuntimeError("Can't find any android device/emulator")
        elif serial is None and len(device_list) > 1:
            raise RuntimeError("more than one device/emulator, please specify the serial number")
        elif serial is None and len(device_list) == 1:
            return device_list[0]
        elif serial is not None and serial in device_list:
            return serial
        raise RuntimeError(f"Can't find android device/emulator {serial}")

    def get_display(self, serial=None):
        """
        获取设备屏幕分辨率
        :param serial: 设备sn号
        :return:
        """
        if serial is None:
            self.must_one_device(serial)
            serial = self.device_list()[0]
        display_str = self.shell(serial, "dumpsys window displays | grep 'app=' | grep 'cur='")
        display = re.findall("cur=([0-9]+)x([0-9]+)", display_str)
        if not display:
            raise ValueError("get screen size failed.")
        return display[0]

    def get_display_rotation(self, serial=None) -> int:
        if serial is None:
            self.must_one_device(serial)
            serial = self.device_list()[0]
        display_str = self.shell(serial, "dumpsys input | grep SurfaceOrientation")
        display = re.findall("[0-9]+", display_str)
        if not display:
            raise ValueError("get display rotation failed.")
        return int(display[0])

    def makedirs(self, serial=None, dirs=None):
        if serial is None:
            self.must_one_device(serial)
            serial = self.device_list()[0]
        msg = self.shell(serial, f"ls {dirs}")
        if "No such file or directory" in msg:
            self.shell(serial, f"mkdir {dirs}")
        return True

    def exists_file(self, serial=None, file=None):
        if serial is None:
            self.must_one_device(serial)
            serial = self.device_list()[0]
        msg = self.shell(serial, f"ls {file}")
        if "No such file or directory" in msg:
            return False
        return True


if __name__ == '__main__':
    adb = AdbClient()
    print([adb.get_display()])
